use thiserror::Error;
use std::collections::HashMap;
use AST::*;
use ParseError::*;

pub type Atom = (String, usize);
pub type Molecule = Vec<Atom>;

#[derive(Error, Debug)]
pub enum ParseError {
    #[error("Not a valid molecule")]
    UnexpectedSymbol(char),
    #[error("Mismatched parenthesis")]
    MismatchedParen,
}
#[derive(Debug)]
pub enum AST {
    Cons(Vec<(AST, usize)>),
    Elem(String),
}
pub fn matched(c1 : char, c2: char) -> bool {
    match (c1,c2) {
        ('(',')') | ('[',']') | ('{','}') => true,
        _ => false,
    }
}
pub fn parse_num(s: &Vec<char>) -> (usize,Vec<char>) {
    let (sn, rest) = read_until(s, 0, |c| !c.is_numeric());
    if sn.len() == 0 {
        return (1, s.clone());
    }
    else {
        return (sn.parse().unwrap(), rest);
    }
}
pub fn read_until<Cond>(s: &Vec<char>, start: usize, cond: Cond) -> (String, Vec<char>) where Cond: FnMut(&char)->bool{
    match s[start..].iter().position(cond) {
        Some(i) => (s[0..(i+start)].iter().collect(), s[(i+start)..].to_vec()),
        None => (s.iter().collect(), vec![]),
    }
}
pub fn read_elem(s: &Vec<char>) -> (String, Vec<char>) {
    return read_until(s, 1, |c| !c.is_lowercase());
}

pub fn parse_until_char(s: &Vec<char>, c: Option<char>) -> Result<(AST, Vec<char>) , ParseError>{
    let mut ss: Vec<char> = s.clone();
    let mut asts = vec![];
    while ss.first().copied() != c {
        if ss.len() == 0 {
            return Err(MismatchedParen);
        }
        let (ast, rest) = parse_atom(ss)?;
        println!("{:?}, {:?}",ast,rest);
        ss = rest;
        if ss.len() > 0 {
            let (n, rest) = parse_num(&ss);
            ss = rest;
            asts.push((ast, n));
        }
        else {
            asts.push((ast, 1));
        }
        println!("asts pushed: {:?}",asts);
    }
    return Ok((Cons(asts), ss));
}
pub fn parse_atom(s: Vec<char>) -> Result<(AST, Vec<char>), ParseError> {
    println!("{:?}",s);
    match s[0] {
        c if c.is_uppercase() => {
            println!("s: {:?}", s);
            let (el, rest) = read_elem(&s);
            println!("el: {:?}", el);
            return Ok((Elem(el), rest));
        },
        '(' | '[' | '{' => {
            let end = match s[0] {
                '(' => ')',
                '[' => ']',
                '{' => '}',
                _   => '?',
            };
            let (mol, rest) = parse_until_char(&s[1..].to_vec(), Some(end))?;
            return Ok((mol, rest[1..].to_vec()));
        },
        _ => {
            return Err(UnexpectedSymbol(s[0]));
        }
    }
}

pub fn traverse(ast: AST, h: &mut HashMap<String, usize>, n: usize) -> Result<(), ParseError> {
    match ast {
        Cons(v) => {
            for t in v {
                let (elem, num) = t;
                match elem {
                    Elem(el) => {
                        *h.entry(el).or_insert(0) += num*n;
                    },
                    _ => {
                        match traverse(elem, h, num*n) {
                            Ok(()) => (),
                            Err(e) => {
                                return Err(e);
                            }
                        }
                    }
                }
            }
        },
        Elem(_) => {
            return Err(UnexpectedSymbol('?'));
        },
    }
    Ok(())
}

pub fn parse_molecule(s: &str) -> Result<Molecule, ParseError> {
    println!("{}",s);
    let chars : Vec<char> = s.chars().collect();
    let (ast, _) = parse_until_char(&chars, None)?;
    println!("ast: {:?}",ast);
    let mut h : HashMap<String, usize> = HashMap::new();
    let ret = traverse(ast, &mut h, 1);
    match ret {
        Ok(()) => Ok(h.into_iter().collect()),
        Err(e) => Err(e),
    }
}
fn main() {
    match parse_molecule("Mg4[ON(SO3)2]2") {
        Ok(v) => println!("{:?}", v),
        Err(e) => println!("{}", e),
    }
}
#[cfg(test)]
mod tests {
    use super::{Molecule, parse_molecule};
    
    macro_rules! assert_parse {
        ($formula:expr, $expected:expr, $name:ident) => {
            #[test]
            fn $name() {
                super::assert_parse($formula, &$expected, "");
            }
        }
    }

    mod molecules {
        assert_parse!("H", [("H",1)], hydrogen);
        assert_parse!("O2", [("O",2)], oxygen);
        assert_parse!("H2O", [("H",2),("O",1)], water);
        assert_parse!("Mg(OH)2", [("Mg",1),("O",2),("H",2)], magnesium_hydroxide);
        assert_parse!("K4[ON(SO3)2]2", [("K",4),("O",14),("N",2),("S",4)], fremys_salt);
        assert_parse!("[(C5H5)Fe(CO)2]2", [("C",14),("H",10),("O",4),("Fe",2)], cyclopentadienyliron_dicarbonyl_dimer);
        assert_parse!("NaSO3(OH)(CH2)5CH3", [("C",6),("H",14),("Na",1),("O",4),("S",1)], hexol_sulphate);
        assert_parse!("Db[Ni]La(Bi24(Rb17)11)", [("Bi",24),("Db",1),("La",1),("Ni",1),("Rb",187)], dbnilabirb);
        assert_parse!("Ta24[Ti2[{Co20}]20]17{Lv}", [("Ta",24),("Ti",34),("Co",6800),("Lv",1)], taticolv);
    }

    #[test]
    fn errors() {
        assert_fail("pie", "Not a valid molecule");
        assert_fail("Mg(OH", "Mismatched parenthesis");
        assert_fail("Mg(OH}2", "Mismatched parenthesis");
    }

    fn assert_fail(formula: &str, msg: &str) {
        let result = parse_molecule(formula);
        assert!(result.is_err(), "expected {} {:?} to fail, got {:?}", msg, formula, result.unwrap());
    }

    fn assert_parse(formula: &str, expected: &[(&str, usize)], _mst: &str) {
        let mut expected = expected.into_iter()
        .map(|&(name, usize)| (name.to_owned(), usize))
        .collect::<Molecule>();
        let result = parse_molecule(formula);
        assert!(result.is_ok(), "expected {:?} to pass, got {:?}", formula, result);
        let mut actual = result.unwrap();
        actual.sort();
        expected.sort();
        assert_eq!(actual, expected);
    }
}
